package com.smh.aiagent.app;

import com.smh.aiagent.advisers.MyLoggerAdvisor;
import com.smh.aiagent.chatmemory.DataBasedChatMemory;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.function.Consumer;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * @author shiminghui
 * @date 2025/5/30 19:18
 * @description: TODO
 */
@Component
@Slf4j
public class LoveApp {

    private final ChatClient chatClient;
    @Value("classpath:/prompts/test_prompt_template.st")
    private org.springframework.core.io.Resource systemResource;

    @Resource
    private VectorStore loveAppVectorStore;
    @Resource
    private Advisor loveAppRAGCloudAdvisor;
    @Resource
    private PgVectorStore pgVectorStore;
    @Resource
    private Advisor loveEmptyContextAdvisor;
    @Resource
    private ToolCallback[] allTools;
    @Resource
    private ChatOptions toolCallingChatOptions;
    @Resource
    private ToolCallingManager toolCallingManager;
    @Resource
    private ToolCallbackProvider toolCallbackProvider;

    public LoveApp(ChatModel dashscopeChatModel, DataBasedChatMemory dataBasedChatMemory) {

        // InMemoryChatMemory chatMemory = new InMemoryChatMemory();
        // FileBasedChatMemory fileBasedChatMemory = new
        // FileBasedChatMemory(System.getProperty("user.dir") +
        // "\\temp\\fileBasedChatMemory");

        ChatClient chatClient = ChatClient.builder(dashscopeChatModel)
                .defaultSystem("你是一个私人助理,你要回答主人的问题")
                .defaultAdvisors(
                        // new MessageChatMemoryAdvisor(chatMemory),
                        // new MessageChatMemoryAdvisor(fileBasedChatMemory),
                        new MessageChatMemoryAdvisor(dataBasedChatMemory), // 这里重新启用,将PG
                        // VectorStore手动创建
                        new MyLoggerAdvisor()
                        // new ReReadingAdvisor()
                )
                .build();

        this.chatClient = chatClient;
    }

    public String doChat(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt(new Prompt(new UserMessage(message)))
                .advisors(new Consumer<ChatClient.AdvisorSpec>() {
                    @Override
                    public void accept(ChatClient.AdvisorSpec advisorSpec) {
                        advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId) // 会话id
                                .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10); // 检索历史记录条数,检索最新的一条
                    }
                }).call().chatResponse();

        return chatResponse.getResult().getOutput().getText();
    }

    /**
     * 将结果返回为结构化的实体
     *
     * @param message
     * @param chatMemoryId
     * @return
     */
    public List<TouristAttractions> doChatWithStruct(String message, String chatMemoryId) {
        List<TouristAttractions> entity = chatClient.prompt(message).advisors(spec -> {
                    spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId) // 会话id
                            .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10);
                }).call()
                // .entity(TouristAttractions.class);
                .entity(new ParameterizedTypeReference<List<TouristAttractions>>() {
                });
        return entity;
    }

    /**
     * 使用字符串模板技术,对系统预设进行模板化处理
     *
     * @param message
     * @param chatMemoryId
     * @return
     */
    public String doChatWithSpringTemplate(String message, String chatMemoryId) {

        // 这里的new Prompt(new UserMessage(message), systemMessage)
        // 顺序必须固定,系统在前,用户在后,最后还是改成api调用吧
        ChatResponse chatResponse = chatClient.prompt()
                .system(spec -> spec.text(systemResource).param("occupational", "程序员"))
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .call().chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    /**
     * 使用QuestionAnswerAdvisor(基于内存的VectorStore)实现简易的RAG
     *
     * @param message
     * @param chatMemoryId
     * @return
     */
    public String doChatWithRAG(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .system("扮演深耕恋爱心理领域的专家。开场向用户表明身份，告知用户可倾诉恋爱难题。围绕单身、恋爱、已婚三种状态提问：单身状态询问社交圈拓展及追求心仪对象的困扰；恋爱状态询问沟通、习惯差异引发的矛盾；已婚状态询问家庭责任与亲属关系处理的问题。引导用户详述事情经过、对方反应及自身想法，以便给出专属解决方案。")
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(new QuestionAnswerAdvisor(loveAppVectorStore))
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    /**
     * 使用dashscope的云知识库检索
     *
     * @param message
     * @param chatMemoryId
     * @return
     */
    public String doChatWithCloudRAG(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .system("扮演深耕恋爱心理领域的专家。开场向用户表明身份，告知用户可倾诉恋爱难题。围绕单身、恋爱、已婚三种状态提问：单身状态询问社交圈拓展及追求心仪对象的困扰；恋爱状态询问沟通、习惯差异引发的矛盾；已婚状态询问家庭责任与亲属关系处理的问题。引导用户详述事情经过、对方反应及自身想法，以便给出专属解决方案。")
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(loveAppRAGCloudAdvisor)
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    /**
     * 使用PGVectorStore进行向量查询
     *
     * @param message
     * @param chatMemoryId
     * @return
     */
    public String doChatWithPGVector(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .system("扮演深耕恋爱心理领域的专家。开场向用户表明身份，告知用户可倾诉恋爱难题。围绕单身、恋爱、已婚三种状态提问：单身状态询问社交圈拓展及追求心仪对象的困扰；恋爱状态询问沟通、习惯差异引发的矛盾；已婚状态询问家庭责任与亲属关系处理的问题。引导用户详述事情经过、对方反应及自身想法，以便给出专属解决方案。")
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(new QuestionAnswerAdvisor(pgVectorStore))
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    /**
     * 添加一个空上下文处理,当检索不到数据时,指定AI如何回答
     *
     * @param message
     * @param chatMemoryId
     * @return
     */
    public String doChatWithEmptyDocument(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(loveEmptyContextAdvisor)
                .call()
                .chatResponse();

        return chatResponse.getResult().getOutput().getText();
    }

    public String doChatWithTools(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .tools(allTools)
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    public String doChatWithToolContext(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .tools(allTools)
                .toolContext(Map.of("userId", 1L))
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    public String doChatWithChatOptions(String message, String chatMemoryId) {

        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .tools(allTools)
                .toolContext(Map.of("userId", 1L))
                .options(toolCallingChatOptions)
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();

    }

    public String doChatWithMcp(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .tools(toolCallbackProvider)
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    /**
     * 使用同步模式进行调用
     *
     * @param message
     * @param chatMemoryId
     * @return
     */
    public String doChatWithSync(String message, String chatMemoryId) {
        ChatResponse chatResponse = chatClient.prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatMemoryId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .tools(allTools)
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    /**
     * 使用流式模式进行调用
     *
     * @param message
     * @param chatId
     * @return
     */
    public Flux<String> doChatByStream(String message, String chatId) {
        return chatClient
                .prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .tools(allTools)
                .stream()
                .content();
    }



    record TouristAttractions(String name, String address, List<String> description) {
    }

}
